import asyncio, socket, struct
from network.package import packHeadLength, JsonPackage, BytePackage, readBinary, magicMap, defaultVer
import logging
from exception import DisconnectedError

logger = logging.getLogger(__name__)

class BaseSocket():

    def __init__(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.loop = asyncio.get_event_loop()

    async def sock_recv(self, sock, nBytes: int=255) -> bytes:
            request = await self.loop.sock_recv(sock, 255)
            if request == b'':
                # sock has been disconnected
                sock.close()
                raise DisconnectedError("Client disconnected")
            return request

    async def receivePackage(self, sock, reqBuf: bytearray):
        bufLen = len(reqBuf)
        request = None
        while bufLen < packHeadLength:
            request = await self.sock_recv(sock)
            bufLen += len(request)
            reqBuf.extend(request)

        (ver, magic, server, plen) = struct.unpack('>BBHI', reqBuf[:packHeadLength])
        while bufLen < plen:
            request = await self.sock_recv(sock)
            bufLen += len(request)
            reqBuf.extend(request)
        pack = readBinary(bytes(reqBuf))
        logger.debug("package received")
        return (pack, reqBuf[plen:])

    async def sendPackage(self, sock, pack):
        logger.debug("sending package")
        await self.loop.sock_sendall(sock, bytes(pack))
        logger.debug("package sent")
    
    def close(self):
        self.sock.close()


class ServerSocket(BaseSocket):
    """
    服务器端 socket 接口
    """

    def __init__(self, host: str="", port: int=8888):
        super(ServerSocket, self).__init__()
        self.sock.bind((host, port))
        self.sock.setblocking(False)
        self.sock.listen(16)
        self.host = (host, port)

    def run(self, controller):
        """
        controller 是处理用户请求的函数

        controller(pack: package, user: dict) -> package
            package 代表 PackageHead, JsonPackage, BytePackage 中的一种
        """
        self.controller = controller
        self.loop.run_until_complete(self.runServer())

    async def runServer(self):
        self.status = True
        while True:
            client, addr = await self.loop.sock_accept(self.sock)
            logger.info("connected: " + str(addr))
            self.loop.create_task(self.handleClient(client))

    async def handleClient(self, clientSock):
        request = None
        reqBuf = bytearray()
        user = None # 记录 socket 连接对应的用户状态
        while self.status:
            try:
                (pack, reqBuf) = await self.receivePackage(clientSock, reqBuf)
            except DisconnectedError as e:
                logger.info("Client disconnected")
                break
            (pack, user) = self.controller(pack, user)
            await self.sendPackage(clientSock, pack)


class ClientSocket(BaseSocket):
    """
    客户端 socket 接口
    """

    def __init__(self, host: str="", port: int=8889):
        super(ClientSocket, self).__init__()
        self.sock.connect((host, port))
        self.reqBuf = bytearray()

    async def sendAsync(self, pack):
        await self.sendPackage(self.sock, pack)
        (result, self.reqBuf) = await self.receivePackage(self.sock, self.reqBuf)
        return result
    
    def send(self, pack):
        """
        发送 pack(JsonPackage/BytePackage) 给服务器
        返回服务器响应结果
        """
        pack = self.loop.run_until_complete(self.sendAsync(pack))
        return pack

    def exit(self):
        """
        断开与服务器的连接
        """
        pack = PackageHead()
        pack._setHeadData(defaultVer, magicMap["disconnect"], 0, packHeadLength)
        self.send(pack)
